import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import xml.etree.ElementTree as ET
from glob import glob
from scipy.ndimage import gaussian_filter
from collections import OrderedDict
import json


def json_points(label):
    with open(label, 'r') as f:
        info = json.load(f)
    points = info['points']
    return points


def xml_points(label):
    tree = ET.parse(label)
    root = tree.getroot()
    points = [[int(obj.find('point').find('x').text), int(obj.find('point').find('y').text)]for obj in root.findall('object')]
    return points


def LMDS_counting(input):
    input_max = torch.max(input).item()

    keep = nn.functional.max_pool2d(input, (3, 3), stride=1, padding=1)
    keep = (keep == input).float()
    input = keep * input

    '''set the pixel valur of local maxima as 1 for counting'''
    input[input < 100.0 / 255.0 * input_max] = 0
    input[input > 0] = 1

    ''' negative sample'''
    if input_max < 0.1:
        input = input * 0

    count = int(torch.sum(input).item())

    return count


def show_comparisons(image_names, comp_path):
    for image_name in image_names:
        image_file = os.path.join(comp_path, str(image_name)+".png")
        img = plt.imread(image_file)
        plt.figure(figsize=(16, 12))
        plt.imshow(img)
        plt.axis('off')
        plt.show()

    




gt_root = '/data/store1/nzd/tir_cc/datasets/rgbtcc/test'

fidtm_results = '/data/store1/nzd/tir_cc/methods/fidtm/25_optim_rgbtcc/results'

gt_files = glob(os.path.join(gt_root, '*_GT.json'))
static = {}
for gt_file in gt_files:
    name = gt_file.split('/')[-1].split('.')[0].replace('_GT', 'R')
    if name in ['2172R', '1213R', '2171R', '2170R']:
        continue
    points = json_points(gt_file)
    gt_num = len(points)

    pre_count = LMDS_counting(torch.load(os.path.join(fidtm_results, f'{name}.pt')))

    static[name] = abs(gt_num - pre_count)
    
static = OrderedDict(sorted(static.items(), key=lambda x: x[1]))



good_names = [k for k, v in static.items()][:200]
bad_names = [k for k, v in static.items()][-200:]



rgb_methods_comp = '02_comp_rgb_rgbtcc'
rgbt_methods_comp = '03_comp_rgbt_rgbtcc'



show_comparisons(good_names, rgb_methods_comp)
show_comparisons(bad_names, rgb_methods_comp)
show_comparisons(good_names, rgbt_methods_comp)
show_comparisons(bad_names, rgbt_methods_comp)
